import numpy as np
import torch

# from . import voxel_interpolation_cpu as voxel_interpolation
import voxel_interpolation_cpu as voxel_interpolation

def segments_intr(a, b, c, d):
    area_abc = (a[0] - c[0]) * (b[1] - c[1]) - (a[1] - c[1]) * (b[0] - c[0])
    area_abd = (a[0] - d[0]) * (b[1] - d[1]) - (a[1] - d[1]) * (b[0] - d[0])
    if (area_abc * area_abd >= 0):
        return None

    area_cda = (c[0] - a[0]) * (d[1] - a[1]) - (c[1] - a[1]) * (d[0] - a[0]);
    area_cdb = area_cda + area_abc - area_abd
    if (area_cda * area_cdb >= 0):
        return None

    t = area_cda / (area_abd - area_abc)
    dx = t * (b[0] - a[0])
    dy = t * (b[1] - a[1])
    return [a[0] + dx ,  a[1] + dy]

def outlier_gt_filter(gt_dense_pts, pts_range):
    lines_num = int(np.max(gt_dense_pts[:, 0]) + 1)

    pts_mask = np.logical_and(gt_dense_pts[:, 1] >= pts_range[0], gt_dense_pts[:, 1] <= pts_range[3])
    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] <= pts_range[4])
    pts_mask = np.logical_and(pts_mask, gt_dense_pts[:, 2] >= pts_range[1])
    new_gt_dense_pts = np.zeros(gt_dense_pts.shape, dtype=np.float32)

    id_count = 0
    filled_size = 0
    for i in range(lines_num):
        pos = (gt_dense_pts[:, 0] == i)
        if (pts_mask[pos].sum() == 0):
            continue

        lines_gt = gt_dense_pts[pos, :]
        lines_gt[:, 0] = id_count
        id_count += 1

        lines_mask = pts_mask[pos]
        n = len(lines_mask)

        if (pts_mask[pos].sum() != pos.sum()):
            for j in range(n-1):
                if (lines_mask[j] != lines_mask[j+1]):
                    p = j
                    break

            inter = [0, 0, 0, 0]
            inter[0] = segments_intr(lines_gt[p, 1:3], lines_gt[p + 1, 1:3], [pts_range[0], pts_range[1]], [pts_range[0], pts_range[4]])
            inter[1] = segments_intr(lines_gt[p, 1:3], lines_gt[p + 1, 1:3], [pts_range[0], pts_range[1]], [pts_range[3], pts_range[1]])
            inter[2] = segments_intr(lines_gt[p, 1:3], lines_gt[p + 1, 1:3], [pts_range[3], pts_range[1]], [pts_range[3], pts_range[4]])
            inter[3] = segments_intr(lines_gt[p, 1:3], lines_gt[p + 1, 1:3], [pts_range[0], pts_range[4]], [pts_range[3], pts_range[4]])
            for item in inter:
                if item is not None:
                    inter = item
                    break

            if lines_mask[p]:
                new_gt_dense_pts[filled_size: filled_size+len(lines_gt[ :p+2, :]), :] = lines_gt[ :p+2, :]
                new_gt_dense_pts[filled_size + p + 1, 1] = inter[0]
                new_gt_dense_pts[filled_size + p + 1, 2] = inter[1]

                filled_size += len(lines_gt[ :p+2, :])
            else:
                new_gt_dense_pts[filled_size:  filled_size + len(lines_gt[ p:, :]), :] = lines_gt[ p:, :]
                new_gt_dense_pts[filled_size, 1] = inter[0]
                new_gt_dense_pts[filled_size, 2] = inter[1]
                filled_size += len(lines_gt[ p:, :])
        else:
            new_gt_dense_pts[filled_size: filled_size + len(lines_gt), :] = lines_gt
            filled_size += len(lines_gt)

    new_gt_dense_pts = new_gt_dense_pts[:filled_size]
    return new_gt_dense_pts

def gt_convert(gt_pts, gt_lines_type):
    pts_num = gt_pts.shape[0]

    new_gt = np.zeros((pts_num, 5))  #[line_id, x, y, z, cls]
    for i in range(pts_num):
        new_gt[i, 0] = gt_pts[i, 0]
        new_gt[i, 1] = gt_pts[i, 1]
        new_gt[i, 2] = gt_pts[i, 2]
        new_gt[i, 3] = 0.0
        new_gt[i, 4] = gt_lines_type[int(gt_pts[i, 0]) , 1]

    return new_gt

if __name__ == '__main__':
    gt_pts = np.array([[0., -72.54324, 5.0472636],
                       [0., -71.49579, 5.0668125],
                       [0., -71.47015, 3.0845041],
                       [0., -71.65209, 3.3092248],
                       [1., -65.04716, 3.1382303],
                       [1., -64.81699, 3.3879757],
                       [1., -64.54222, 4.1366506],
                       [1., -64.35778, 4.8423443],
                       [1., -57.427032, 4.923033],
                       [1., -57.480003, 4.005681],
                       [1., -57.40822, 3.228975],
                       [1., -57.192364, 2.8135931],
                       [2., -50.98361, 3.026371],
                       [2., -50.82041, 3.2991142],
                       [2., -50.60953, 4.0911403],
                       [2., -50.578136, 4.991172],
                       [2., -43.345768, 4.7872505],
                       [2., -43.258366, 4.650796],
                       [2., -43.259262, 3.7984395],
                       [2., -43.01894, 3.3023226],
                       [2., -42.73904, 3.0122755],
                       [3., -30.559416, 2.5667028],
                       [3., -30.145872, 3.0045815],
                       [3., -29.680273, 3.684008],
                       [3., -29.55969, 4.173395],
                       [3., -29.359497, 5.008037],
                       [3., -26.252169, 4.888068],
                       [3., -26.087988, 3.275507],
                       [3., -25.829792, 2.7363052],
                       [4., 17.54122, 3.6259096],
                       [4., 17.900475, 4.2746634],
                       [4., 23.332075, 4.4806542],
                       [4., 23.441315, 3.8365023],
                       [4., 23.506542, 3.6869116],
                       [5., 40.01681, 2.2263389],
                       [5., 39.938576, 3.1677864],
                       [5., 39.91483, 3.9010007],
                       [5., 40.966946, 3.9204943],
                       [5., 41.05006, 18.135746],
                       [5., 45.46731, 18.199795],
                       [5., 45.570663, 2.8264313],
                       [5., 45.609463, 2.3562086],
                       [5., 45.69426, 1.9928378],
                       [6., 19.069717, -17.878555],
                       [6., 19.071545, -18.452492],
                       [6., 19.185635, -22.653385],
                       [6., 12.263296, -22.66391],
                       [6., 12.032299, -17.232298],
                       [6., 15.505721, -17.823833],
                       [7., -38.58253, -25.432722],
                       [7., -38.636036, -27.221443],
                       [7., -38.46231, -29.966648],
                       [7., -45.139435, -30.499338],
                       [7., -44.71966, -27.318247],
                       [7., -44.684437, -25.088486],
                       [7., -41.679367, -25.389835],
                       [8., -73.28628, -5.3868575],
                       [8., -73.388275, -5.4761705],
                       [8., -73.71425, -6.449335],
                       [8., -73.76602, -7.146481],
                       [8., -78.66249, -7.2136345],
                       [8., -78.80498, -6.526554],
                       [8., -78.97091, -5.7215843],
                       [8., -79.10881, -5.525136],
                       [9., -77.820595, -1.0803046],
                       [9., -43.43571, -1.4397023],
                       [10., -41.212147, -6.67708],
                       [10., -41.679367, -25.389835],
                       [11., -12.046524, -1.8275703],
                       [11., 2.7842429, -1.8646312],
                       [12., 26.81533, -2.051773],
                       [12., 38.503353, -2.14281],
                       [13., 15.334649, -9.848127],
                       [13., 15.505721, -17.823833]], dtype=np.float32)

    gt_lines_type = np.array([[0, 0],
                              [1, 0],
                              [2, 0],
                              [3, 0],
                              [4, 0],
                              [5, 0],
                              [6, 0],
                              [7, 0],
                              [8, 0],
                              [9, 1],
                              [10, 1],
                              [11, 1],
                              [12, 1],
                              [13, 1]], dtype=np.int32)

    voxels_size = np.array([0.8, 0.8, 80.], dtype=np.float32)
    pts_range = np.array([-51.2, -51.2, -10., 51.2, 51.2, 10.], dtype=np.float32)

    gt_dense_pts = gt_convert(gt_pts, gt_lines_type)
    gt_dense_pts = outlier_gt_filter(gt_dense_pts, pts_range)
    gt_dense_pts = gt_dense_pts.astype(np.float32)
    k_max = 10000
    cur_batch_dense_pts = np.zeros((k_max, 8), dtype=np.float32)   # [line_id, pos, x, y, z, cw_distance, ccw_distance, cls]

    size = voxel_interpolation.gt_lines_voxel_interpolation(gt_dense_pts, voxels_size, pts_range, k_max, cur_batch_dense_pts)
    cur_batch_dense_pts = cur_batch_dense_pts[:size, :]
    print(size)



    out_dict = {
        'cur_batch_dense_pts': cur_batch_dense_pts,
        'gt_pts': gt_pts,
        'gt_lines_type': gt_lines_type,
        'gt_dense_pts': gt_dense_pts,
    }
    output = open('/home/zlf/tmp_decode/2.pkl', 'wb')
    import pickle
    pickle.dump(out_dict, output)
    output.close()